import cv2
from cv2 import dnn_superres
import numpy as np
import os.path as osp

BASEDIR = osp.abspath(osp.dirname(__file__))

# 创建SR对象...
sr = dnn_superres.DnnSuperResImpl_create()
# 读入模型
path_model = osp.join(BASEDIR, "SR_weights", "EDSR_x4.pb")
sr.readModel(path_model)
# 设定算法和放大比例
sr.setModel("edsr", 4)


def apply_single_EDSR(arr_img, size_back=True):
    """
    输入一张图片的ndarray，输出超分辨后的ndarray。
    """
    result_img = sr.upsample(arr_img)
    # 重新压缩回来
    if size_back:
        h, w, c = arr_img.shape
        result_img = cv2.resize(result_img, (w, h), interpolation=cv2.INTER_CUBIC)
    return result_img


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    
    # img_input_path = 'D:/TZB/plane_reid_addTriplet1/data/train_planes_cut/D_pic389_21.png'
    # img_input_path = 'D:/TZB/plane_reid_addTriplet1/data/train_planes_cut/C_pic6_8.png'
    img_input_path = 'D:/TZB/plane_reid_addTriplet1/data/train_planes_cut/C_pic110_0.png'
    
    img = cv2.imread(img_input_path)
    print(img.shape)
    
    h, w, c = img.shape    
    if h < 64 or w < 64:
        result_img = apply_single_EDSR(input_img_array=img, size_back=False)
        stack_img = np.hstack((cv2.resize(img, (img.shape[1]*4, img.shape[0]*4), interpolation=cv2.INTER_CUBIC), result_img))
    else:
        result_img = apply_single_EDSR(input_img_array=img)
        stack_img = np.hstack((img, result_img))
        
    print(result_img.shape)
    plt.imshow(stack_img)
    plt.show()
    